#
#  MolViz.py
#  PyMol Plugin
#
#  Created by Christian Muise on 25/01/08.
#  Copyright (c) 2008 Christian Muise. All rights reserved.
#

from pymol import cmd
from time import sleep
from SocketServer import Server
import re, math, time


class MolViz:
	def __init__(self):
	
		self.message_pattern = re.compile(r'X(.+);Y(.+);Z(.+);')
		self.change_pattern = re.compile(r'(.+):(.+)')
		
		#----- Vars that are updated and used to turn things
		self.x_angle = 0
		self.y_angle = 0
		self.x = 0
		self.y = 0
		self.z = 30
		
		#--- To prevent nasty race conditions
		self.x_tmp = 0
		self.y_tmp = 0
		self.z_tmp = 0
		self.lock = 0
		
		self.old_zoom = 0
		
		self.frame_rate_in = 0
		self.frame_rate_out = 0
		
		self.zoom_error_count = 0
		self.zoom_error_max = 3
		self.zoom_moving = -1
		self.zoom_error_threshold = 0.03
		
		#----- Settings
		#--- Screen ratio defines how close the screen is to the user in relation to the origin
		self.screen_ratio = 0.5
		#--- Screen origin is if we assume the user is looking at the screen or behind it.
		self.screen_origin = 0
		
		#--- Accelerated rotation
		self.do_accelerated_rotation = -1    # -1 means don't do it and 1 means do it
		self.accelerated_rotation_angle = 90 # what we should be able to rotate to
		self.accelerated_rotation_angle_rad = math.radians(90) # what we should be able to rotate to
		self.accelerated_rotation_fov = 40   # what the current field of view is
		
		#----- Setup for the running average
		self.running_avg_index_depth = 0
		self.running_avg_weight_depth = 1
		
		self.running_avg_index_rot = 0
		self.running_avg_weight_rot = 4
		
		self.running_avg_x = 0
		self.running_avg_nums_x = [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
		
		self.running_avg_y = 0
		self.running_avg_nums_y = [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
		
		self.running_avg_z = 0
		self.running_avg_nums_z = [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
		
		#----- The current view matrix
		self.current_view = cmd.get_view()
		
		#----- The server used to listen for head tracking
		self.server = Server(self, 4440)
		
		#----- The server used to listen for the GUI to update the values / display info
		self.gui = Server(self, 4441)
		
		#----- Initial value guesses
		self.old_zoom = self.current_view[11]
		self.ratio = self.old_zoom / self.z
		
	
	
	###################
	##   Gets/Sets   ##
	###################
	
	def set_data(self, data):
		result = data.rstrip(' ')
		result = result.rstrip('\n')
		if (result != ''):
			(x,y,z) = self.message_pattern.search(result).groups()
			if self.lock != 1:
				self.x_tmp = float(x)
				self.y_tmp = float(y)
				self.z_tmp = float(z)
	
	def set_change(self, change):
		result = change.rstrip(' ')
		result = result.rstrip('\n')
		if (result != ''):
			
			if re.match('togAccRot', result):
				self.toggleAcceleratedRotation()
			elif re.match('setAccRot:(.+)', result):
				(cmd, rot) = self.change_pattern.search(result).groups()
				self.setAcceleratedRotation(rot)
			elif re.match('setAccFov:(.+)', result):
				(cmd, fov) = self.change_pattern.search(result).groups()
				self.setAcceleratedRotationFov(fov)
			elif re.match('setRotWt:(.+)', result):
				(cmd, wt) = self.change_pattern.search(result).groups()
				self.setRotationWeight(wt)
			elif re.match('setZmWt:(.+)', result):
				(cmd, wt) = self.change_pattern.search(result).groups()
				self.setZoomWeight(wt)
			elif re.match('setZmRat:(.+)', result):
				(cmd, rat) = self.change_pattern.search(result).groups()
				self.setZoomRatio(rat)
			elif re.match('setMaxErr:(.+)', result):
				(cmd, max) = self.change_pattern.search(result).groups()
				self.setMaxError(max)
			elif re.match('setErrThresh:(.+)', result):
				(cmd, thresh) = self.change_pattern.search(result).groups()
				self.setErrorThreshold(thresh)
			elif re.match('setScrPlc:(.+)', result):
				(cmd, plc) = self.change_pattern.search(result).groups()
				self.setScreenPlacement(plc)
			elif re.match('setScrPlcRat:(.+)', result):
				(cmd, rat) = self.change_pattern.search(result).groups()
				self.setScreenPlacementRatio(rat)
			elif re.match('autoCV', result):
				self.calibrate('cv')
			elif re.match('autoIR', result):
				self.calibrate('ir')
			else:
				print "Bad message: " + str(result)
			
	def set_fps_in(self, fps):
		self.frame_rate_in = float(fps)
	
	def toggleAcceleratedRotation(self):
		self.do_accelerated_rotation *= -1
		
	def setAcceleratedRotation(self, rot):
		self.accelerated_rotation_angle = float(rot)
		self.accelerated_rotation_angle_rad = math.radians(self.accelerated_rotation_angle)
		
	def setAcceleratedRotationFov(self, fov):
		self.accelerated_rotation_fov = float(fov)
	
	def setRotationWeight(self, w):
		self.running_avg_weight_rot = int(w)
	
	def setZoomWeight(self, w):
		self.running_avg_weight_depth = int(w)
	
	def setZoomRatio(self, r):
		self.ratio = float(r)
	
	def setMaxError(self, m):
		self.zoom_error_max = int(m)
	
	def setErrorThreshold(self, t):
		self.zoom_error_threshold = float(t)
	
	def setScreenPlacement(self, p):
		self.screen_origin = int(p)
		if self.screen_origin == 1:
			cmd.zoom()
			cmd.translate([0,0,30])
			cmd.clip("near", 30)
			
		elif self.screen_origin == 0:
			cmd.orient()
			cmd.zoom()
	
	def setScreenPlacementRatio(self, r):
		self.screen_ratio = float(r)
	
	
	
	
	def run(self):
	
		self.server.start()
		self.gui.start()
	
		running = 1
		frame_count = 0
		timestamp = time.time()
		
		while running:
			self.current_view = cmd.get_view()
			
			#--- SAFELY get the x,y,z
			while self.lock == 1:
				sleep(0.01)
			self.lock = 1
			self.x = self.x_tmp
			self.y = self.y_tmp
			self.z = self.z_tmp
			self.lock = 0
			
			new_zoom = self.z * self.ratio
			
			#----- Rotation
			frac_x = self.x / (0.5 * math.radians(self.accelerated_rotation_fov))
			frac_y = self.y / (0.375 * math.radians(self.accelerated_rotation_fov))
			
			if abs(frac_x) > 1:
				self.x = self.running_avg_x
			else:
				#--- See if we want to do accelerated rotation
				if self.do_accelerated_rotation > 0:
					# math.radians = 1.5707963267948966
					self.x = (math.asin(frac_x) / 1.5707963267948966) * self.accelerated_rotation_angle_rad
			
			if abs(frac_y) > 1:
				self.y = self.running_avg_y
			else:
				#--- See if we want to do accelerated rotation
				if self.do_accelerated_rotation > 0:
					self.y = (math.asin(frac_y) / 1.5707963267948966) * self.accelerated_rotation_angle_rad
			
			#self.x = (float(self.accelerated_rotation_angle) / float(self.accelerated_rotation_fov)) * self.x
				
				
			#--- Safety mechanism for rotating.
			if (abs(math.degrees(self.x)) > 175):
				self.x = self.running_avg_x
			if (abs(math.degrees(self.y)) > 175):
				self.y = self.running_avg_y
			
			
			for i in range(self.running_avg_weight_rot):
				self.running_avg_x -= self.running_avg_nums_x[self.running_avg_index_rot]
				self.running_avg_nums_x[self.running_avg_index_rot] = (self.x / 20)
				self.running_avg_x += self.running_avg_nums_x[self.running_avg_index_rot]
				
				self.running_avg_y -= self.running_avg_nums_y[self.running_avg_index_rot]
				self.running_avg_nums_y[self.running_avg_index_rot] = (self.y / 20)
				self.running_avg_y += self.running_avg_nums_y[self.running_avg_index_rot]
				
				self.running_avg_index_rot = (self.running_avg_index_rot + 1) % 20
			
			
			#----- Zoom
			zoom_error = abs(abs(new_zoom) - abs(self.running_avg_z)) / (abs(self.running_avg_z) + 0.00000000001)
			
			if ((self.zoom_moving > 0) and (self.zoom_error_count > 3 * self.zoom_error_max)):
				self.zoom_error_count = 0
				self.zoom_moving = -1
				
			if ((self.zoom_moving < 0) and (self.zoom_error_count > self.zoom_error_max)):
				self.zoom_error_count = 0
				self.zoom_moving = 1
			
			if (zoom_error > self.zoom_error_threshold):
				if (self.zoom_moving > 0):
					self.zoom_error_count = -2 * self.zoom_error_threshold
					for i in range(self.running_avg_weight_depth):
						self.running_avg_z -= self.running_avg_nums_z[self.running_avg_index_depth]
						self.running_avg_nums_z[self.running_avg_index_depth] = (new_zoom / 30)
						self.running_avg_z += self.running_avg_nums_z[self.running_avg_index_depth]
						
						self.running_avg_index_depth = (self.running_avg_index_depth + 1) % 30
				
				else:
					self.zoom_error_count += 1
			
			else:
				if (self.zoom_moving < 0):
					self.zoom_error_count = 0
				
				else:
					self.zoom_error_count += 1
			
			
			diff_x = self.running_avg_x - self.x_angle
			diff_y = self.running_avg_y - self.y_angle
			diff_zoom = self.running_avg_z - self.old_zoom
			
			if (self.screen_origin == 0):
				cmd.turn('y', math.degrees(diff_x))
				cmd.turn('x', math.degrees(diff_y))
			
			elif (self.screen_origin == 1):
				cmd.turn('y', math.degrees(math.atan(self.screen_ratio * math.tan(diff_x))))
				cmd.turn('x', math.degrees(math.atan(self.screen_ratio * math.tan(diff_y))))
			
			else:
				print "Error - Invalid screen mode, " + str(self.screen_origin)
			
			cmd.move('z', diff_zoom)
			
			self.old_zoom = self.running_avg_z
			
			self.x_angle = self.running_avg_x
			self.y_angle = self.running_avg_y
			
			#print "Z (" + str(self.running_avg_z) + ")"
			
			sleep(0.03)
			
			frame_count += 1
			
			if frame_count == 50:
				frame_count = 0
				new_time = time.time()
				fps = 50 / (new_time - timestamp)
				self.frame_rate_out = fps
				timestamp = new_time
				
				print "FPS:" + str(self.frame_rate_in) + ";" + str(self.frame_rate_out) + ";"
	
	def calibrate(self, mode = 'ir'):
		print "Calibrating in 3..."
		sleep(1)
		print "2.."
		sleep(1)
		print "1."
		sleep(1)
	
		cmd.orient()
		cmd.zoom()
		self.current_view = cmd.get_view()
		
		# Sanity check
		if self.z == 0:
			self.ratio = self.current_view[11]
		else:
			self.ratio = self.current_view[11] / self.z
			
		self.old_zoom = self.current_view[11]
		self.zoom_cap = abs(0.001 * self.current_view[11])
		self.zoom_error_count = 0
		
		# mode = 0 -> head tracking with open cv (need huge z stabilization)
		if (mode == 'cv'):
			self.running_avg_weight_depth = 2
			self.running_avg_weight_rot = 4
			self.screen_origin = 0
			
			self.do_accelerated_rotation = -1
			self.accelerated_rotation_angle = 80
			self.accelerated_rotation_angle_rad = math.radians(80)
			self.accelerated_rotation_fov = 50
			
			self.zoom_error_count = 0
			self.zoom_error_max = 5
			self.zoom_moving = -1
			self.zoom_error_threshold = 0.05
		
			print "Mode set to do open cv..."
		
		elif (mode == 'ir'):
			self.running_avg_weight_depth = 3
			self.running_avg_weight_rot = 20
			self.screen_origin = 0
			
			self.do_accelerated_rotation = -1
			self.accelerated_rotation_angle = 90
			self.accelerated_rotation_angle_rad = math.radians(90)
			self.accelerated_rotation_fov = 40
			
			self.zoom_error_count = 0
			self.zoom_error_max = 2
			self.zoom_moving = -1
			self.zoom_error_threshold = 0.02
			
			print "Mode set to do ir tracking..."
			
		else:
			print "Invalid Input Mode - " + mode
		
		# Dump for the gui
		print "UPDATE:" + str(self.accelerated_rotation_angle) + ";" + str(self.accelerated_rotation_fov) + ";" + str(self.running_avg_weight_rot) + ";" + str(self.running_avg_weight_depth) + ";" + str(self.ratio) + ";" + str(self.zoom_error_max) + ";" + str(self.zoom_error_threshold) + ";" + str(self.screen_origin) + ";" + str(self.screen_ratio) + ";"



molviz = MolViz()

def calibrate(param = 'ir'):
	molviz.calibrate(param)

cmd.extend('mv_calibrate', calibrate)

molviz.run()

